{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "source": [
    "!pip install paddle2onnx onnxruntime onnxslim onnxruntime-tools onnx  -i https://pypi.tuna.tsinghua.edu.cn/simple"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "is_executing": true,
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "!paddle2onnx --model_dir rapid_table_det_paddle/models/obj_det --model_filename model.pdmodel --params_filename model.pdiparams --save_file rapid_table_det/models/obj_det.onnx --opset_version 16 --enable_onnx_checker True\n",
    "!paddle2onnx --model_dir rapid_table_det_paddle/models/db_net --model_filename model.pdmodel --params_filename model.pdiparams --save_file rapid_table_det/models/edge_det.onnx --opset_version 16 --enable_onnx_checker True\n",
    "!paddle2onnx --model_dir rapid_table_det_paddle/models/pplcnet --model_filename model.pdmodel --params_filename model.pdiparams --save_file rapid_table_det/models/cls_det.onnx --opset_version 16 --enable_onnx_checker True\n",
    "\n",
    "!onnxslim rapid_table_det/models/obj_det.onnx rapid_table_det/models/obj_det.onnx\n",
    "!onnxslim rapid_table_det/models/edge_det.onnx rapid_table_det/models/edge_det.onnx\n",
    "!onnxslim rapid_table_det/models/cls_det.onnx rapid_table_det/models/cls_det.onnx"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-15T12:12:35.265576Z",
     "start_time": "2024-10-15T12:12:34.281134Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "from pathlib import Path\n",
    "import onnx\n",
    "from onnxruntime.quantization import quantize_dynamic, QuantType, quant_pre_process\n",
    "def quantize_model(root_dir_str, model_dir, pre_fix):\n",
    "\n",
    "    original_model_path = f\"{pre_fix}.onnx\"\n",
    "    # quantized_model_path = f\"{pre_fix}_quantized.onnx\"\n",
    "    quantized_model_path = original_model_path\n",
    "    original_model_path = f\"{root_dir_str}/{model_dir}/{original_model_path}\"\n",
    "    quantized_model_path = f\"{root_dir_str}/{model_dir}/{quantized_model_path}\"\n",
    "    quant_pre_process(original_model_path, quantized_model_path, auto_merge=True)\n",
    "    # 进行动态量化\n",
    "    quantize_dynamic(\n",
    "        model_input=quantized_model_path,\n",
    "        model_output=quantized_model_path,\n",
    "        weight_type=QuantType.QUInt8\n",
    "    )\n",
    "\n",
    "    # 验证量化后的模型\n",
    "    quantized_model = onnx.load(quantized_model_path)\n",
    "    onnx.checker.check_model(quantized_model)\n",
    "    print(\"Quantized model is valid.\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "root_dir_str = \".\"\n",
    "model_dir = f\"rapid_table_det/models\"\n",
    "quantize_model(root_dir_str, model_dir,  \"obj_det\")\n",
    "quantize_model(root_dir_str, model_dir,  \"edge_det\")\n",
    "# quantize_model(root_dir_str, model_dir,  \"cls_det\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [],
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
